#include "tbf.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <sys/time.h>

#define MAXTBF 1024

typedef struct tbf_st {
	int token; // 令牌
	int cps;   // 速率
	int burst; // 上限
} tbf_t;

static tbf_t *libs[MAXTBF];
static int inited;

static void sig_handler(int s)
{
	int td;
	tbf_t *cur;

	for (td = 0; td < MAXTBF; td++) {
		cur = libs[td];
		if (NULL != cur) {
			cur->token += cur->cps;

			if (cur->token > cur->burst) {
				cur->token = cur->burst;
			}
		}
	}
}

static void module_load()
{
	struct sigaction act = {};
	struct sigaction oldact;
	struct itimerval itv, olditv;

	// 设置定时调用函数
	act.sa_handler = sig_handler;
	sigaction(SIGALRM, &act, &oldact);
	itv.it_interval.tv_sec = 0;
	itv.it_interval.tv_usec = 10000; // 0.1s
	itv.it_value.tv_sec = 0;
	itv.it_value.tv_usec = 10000; // 0.1s
	setitimer(ITIMER_REAL, &itv, &olditv);
}

int tbf_init(int cps, int burst)
{
	int td;
	tbf_t *new_tbf; // 创建令牌桶

	for (td = 0; td < MAXTBF; td++) {
		if (NULL == libs[td]) {
			break;
		}
	}
	// 开辟存储空间
	new_tbf = malloc(sizeof (tbf_t));
	// 成员初始化
	new_tbf->token = 0;
	new_tbf->cps = cps;
	new_tbf->burst = burst;
	// 存放到令牌桶中
	libs[td] = new_tbf;
	if (!inited) {
		module_load();
		inited = 1;
	}
	return td; // 返回初始化的桶下标
}

// 返回取到令牌的个数
int tbf_fetch_token(int td, int ntoken)
{
	if (NULL == libs[td])
		return 0;

	while (libs[td]->token < ntoken) { // 现在的token还不够取的，等一会
		pause();
	}
	// 把token取出来，token数减少
	libs[td]->token -= ntoken;
	
	return ntoken; // 取出来这么多个
}

// 返回是否成功
int tbf_return_token(int td, int ntoken)
{
	if (NULL == libs[td])
		return -1;
	// 把token放进去，token数增加
	libs[td]->token += ntoken;

	return 0;
}

void tbf_destroy(int td)
{
	if (NULL != libs[td]) {
		free(libs[td]);

		libs[td] = NULL;
	}
}

void tbf_destroy_all(void)
{
	int td;

	for (td = 0; td < MAXTBF; td++) {
		if (NULL != libs[td]) {
			free(libs[td]);

			libs[td] = NULL;
		}
	}
}

